{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b4f813da",
   "metadata": {},
   "source": [
    "# The `Ramsey` notebook  <a id=\"Ramsey\"></a>[<font size=1>(back to `Main.ipynb`)</font>](./Main.ipynb)\n",
    "\n",
    "This notebook gathers the computations of the steady-state Ramsey solution for the truncated model. \n",
    "\n",
    "Thz computation mostly follows the matrix computation of the [paper](https://francois-le-grand.fr/docs/research/LR_optimal-debt.pdf).\n",
    "\n",
    "To sum it up, we aim at simulating the Ramsey model using perturbation methods via Dynare. To do so, we need to determine the steady-state Ramsey allocation. Our computational strategy consists in finding the Pareto weights of the SWF such that the observed US tax system is optimal at the steady state. It means that the truncated allocation `truncatedModel::TruncatedModel` computed by the function `TruncatedModel` of the notebook [Truncation.ipynb](./Truncation.ipynb) yields the Ramsey allocation (because it has been computed for the observed US fiscal system). However, Ramsey FOCs depend on the Ramsey Lagrange multipliers as well as the Pareto weights and the marginal value of liquidity $(\\psi_{h})_h$. We thus need to compute the steady-state values of these quantities.\n",
    "\n",
    "We do this in two steps. [First](#Ramsey-LM), we express Lagrange multipliers and marginal value of liquidity as a function of Pareto weights only using the truncated allocation `truncatedModel::TruncatedModel`. [Second](#Ramsey-weights), we compute the Pareto weights that respect the Ramsey FOCs, such that the SWF is the closest to the utilitarian welfare. \n",
    "\n",
    "Finally, we present the [implementation](#Ramsey-implementation)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa6e8cec",
   "metadata": {},
   "source": [
    "## The Lagrange multipliers <a id=\"Ramsey-LM\"></a>[<font size=1>(back to Ramsey)</font>](#Ramsey)\n",
    "\n",
    "Computing the Lagrange multipliers as a function of the Pareto weights using the truncated allocation is not difficult but involves some tedious algebra. However, we take advantage of the finite state space that makes it possible to express these Lagrange multipliers (e.g., $(\\lambda_{c,h})_h$ for the Euler equations) as vectors.  The resulting implementation is fast as it mostly involves basic linear algebra. The matrix representation consists in stacking together the equations characterizing the steady state, so as to provide a convenient matrix notation for solving the steady state.\n",
    "\n",
    "\n",
    "### Preliminary definitions \n",
    "\n",
    "We start with a number of definitions that will allow us to express Lagrange multipliers as function of Pareto weight using the truncated allocation: `truncatedModel::TruncatedModel` of the notebook [Truncation.ipynb](./Truncation.ipynb).\n",
    "\n",
    "We denote with a bold letter the vector associated to a given variable. For instance $\\boldsymbol{c} := (c_h)_{h\\in\\mathcal H}$ corresponds to the vector of individual consumption levels in truncated histories $h$. The order in which the histories are stacked does not matter but the order has to be consistent from one vector (or matrix) to another. \n",
    "\n",
    "We now define a number of quantities. \n",
    "\n",
    "* We define the matrix: $\\bar{\\boldsymbol{\\Pi}} = \\boldsymbol{S} \\circ \\boldsymbol{\\Pi}^\\top \\circ (1/\\boldsymbol{S})$, where $\\boldsymbol{S}$ is the size (as a vector) of histories (variable `S_h`) and $\\Pi^\\top$ the transpose of the transition matrix for histories (variable `Π_h` -- be careful that this matrix is actually the transpose of $\\boldsymbol{\\Pi}$).\n",
    "* We define the diagonal matrix $\\boldsymbol{P}$ (`P`) with $1$ on the diagonal if history $h$ is not constrained and $0$ otherwise. \n",
    "* The matrix $\\boldsymbol{D}_\\boldsymbol{x}$ is the diagonal matrix with $\\boldsymbol{x}$ on the diagonal. Note that $\\boldsymbol{D}_\\boldsymbol{x}^{-1} = D_{1/\\boldsymbol{x}}$ if no element of $\\boldsymbol{x}$ is null.\n",
    "\n",
    "\n",
    "Using this notation, we define a series of matrices and vectors:\n",
    "\n",
    "* First:\n",
    "\\begin{align*} \n",
    "\\boldsymbol{M}_{1}&:=\\boldsymbol{P}(\\boldsymbol{I}-\\beta(1+r)\\bar{\\boldsymbol{\\Pi}})\\boldsymbol{D}_{\\boldsymbol{\\xi}^{E}\\circ u^{\\prime\\prime}(\\boldsymbol{x})}(\\boldsymbol{I}-(1+r)\\boldsymbol{\\Pi}^{\\top})+\\boldsymbol{I}-\\boldsymbol{P},\\\\\n",
    "\\boldsymbol{M}_{2}&:=\\boldsymbol{M}_{1}^{-1}\\boldsymbol{P}(\\boldsymbol{I}-\\beta(1+r)\\bar{\\boldsymbol{\\Pi}})\\boldsymbol{D}_{\\boldsymbol{\\xi}^{E}\\circ u^{\\prime}(\\boldsymbol{x})}, \\\\\n",
    "\\boldsymbol{V}_{1}&:=-\\boldsymbol{M}_{1}^{-1}\\boldsymbol{P}(\\boldsymbol{I}-\\beta(1+r)\\bar{\\boldsymbol{\\Pi}})\\boldsymbol{S},\n",
    "\\end{align*} \n",
    "\n",
    "which are represented by the matrices `M1` and `M2` and the vector `V1`, respectively.  Note that we use the left division (with `\\`) rather than the computation of the inverse. Indeed, `A\\B` is more stable and faster than `inv(A)*B` -- but both yield the same result when `A` is invertible.\n",
    "\n",
    "* Second: \n",
    "\\begin{align*} \n",
    "\\boldsymbol{M}_{3}&:=\\boldsymbol{D}_{\\boldsymbol{\\xi}^{E}\\circ u^{\\prime}(\\boldsymbol{x})}-\\boldsymbol{D}_{\\boldsymbol{\\xi}^{E}\\circ u^{\\prime\\prime}(\\boldsymbol{x})}(\\boldsymbol{I}-(1+r)\\boldsymbol{\\boldsymbol{\\Pi}^{\\top}})\\boldsymbol{M}_{2},\\\\\n",
    "\\boldsymbol{V}_{2}&:=-\\boldsymbol{D}_{\\boldsymbol{\\xi}^{E}\\circ u^{\\prime\\prime}(\\boldsymbol{x})}(\\boldsymbol{I}-(1+r)\\boldsymbol{\\boldsymbol{\\Pi}^{\\top}}\\boldsymbol{V}_{1})-\\boldsymbol{S},\n",
    "\\end{align*} \n",
    "\n",
    "which are represented by the matrix `M3`, and the vector `V2`.\n",
    "* Third:\n",
    "\\begin{align*}\n",
    "C_{1}&:=\\boldsymbol{S}^{\\top}(\\frac{l^{1/\\varphi}}{\\chi}\\boldsymbol{y}^{\\tilde{\\tau}}-F_{L}\\boldsymbol{y}^{1+\\frac{\\tilde{\\tau}}{1/\\varphi+1}})-\\frac{1}{\\chi\\tilde{\\tau}}(1/\\varphi+1)l^{1/\\varphi}(\\boldsymbol{y}^{\\tilde{\\tau}})^{\\top}\\boldsymbol{V}_{2},\\\\\n",
    "\\boldsymbol{L}_{0}^{\\top}&:=\\frac{1+1/\\varphi}{\\chi\\tilde{\\tau}}l^{1/\\varphi}(\\boldsymbol{y}^{\\tilde{\\tau}})^{\\top}\\boldsymbol{M}_{3}/C_{1},\n",
    "\\end{align*}\n",
    "\n",
    "which are represented by the scalar `C1` and the vector `L0`.\n",
    "\n",
    "* Fourth: \n",
    "\\begin{align*}\n",
    "\\boldsymbol{M}_{4}&:=\\boldsymbol{M}_{3}+\\boldsymbol{V}_{2}\\boldsymbol{L}_{0}^{\\top} \\\\\n",
    "\\boldsymbol{M}_{5}&:=\\boldsymbol{M}_{2}+\\boldsymbol{V}_{1}\\boldsymbol{L}_{0}^{\\top}\n",
    "\\end{align*}\n",
    "\n",
    "which are represented by the matrices `M4` and `M5`.\n",
    "\n",
    "* Fifth:\n",
    "\\begin{align*}\n",
    "\\boldsymbol{L}_{1}^{\\top}&:=\\left(\\tilde{\\boldsymbol{a}}^{\\top}\\boldsymbol{M}_{4}+\\text{where: }(\\boldsymbol{\\xi}^{E}\\circ u^{\\prime}(\\boldsymbol{x}))^{\\top}\\boldsymbol{\\Pi}^{\\top}\\boldsymbol{M}_{5}\\right),\\\\\n",
    "\\boldsymbol{L}_{2}^{\\top}&:=\\frac{l^{1/\\varphi+1}}{\\chi\\tilde{\\tau}}(\\boldsymbol{y}^{\\tilde{\\tau}}\\!\\circ\\!(\\log\\boldsymbol{y}\\!-\\!\\frac{1}{\\tilde{\\tau}}\\mathbf{1}_{N_{tot}}))^{\\top}\\!\\boldsymbol{M}_{4}\\!-\\!\\frac{l(\\boldsymbol{S}\\!\\circ\\!\\log\\boldsymbol{y})^{\\top}}{1/\\varphi+1}(\\frac{l^{1/\\varphi}}{\\chi}\\boldsymbol{y}^{\\tilde{\\tau}}\\!-\\!F_{L}\\boldsymbol{y}^{\\frac{1/\\varphi+1+\\tilde{\\tau}}{1/\\varphi+1}})\\boldsymbol{L}_{0}^{\\top},\n",
    "\\end{align*}\n",
    "\n",
    "which are represented by the the vectors `L1` and `L2`.\n",
    "\n",
    "* Finally:\n",
    "\\begin{align*}\n",
    "\\boldsymbol{R}_{0}&:=\\left[\\begin{array}{cccc}\n",
    "\\boldsymbol{1}_{N_{tot}^{1}} & 0 & 0 & 0\\\\\n",
    "0 & \\boldsymbol{1}_{N_{tot}^{2}} & 0 & 0\\\\\n",
    "\\vdots & \\vdots & \\ddots & \\vdots\\\\\n",
    "0 & 0 & \\ldots & \\boldsymbol{1}_{N_{tot}^{F}}\n",
    "\\end{array}\\right],\\\\\n",
    "\\boldsymbol{R}_{1}&:=\\boldsymbol{D}_{\\boldsymbol{S}}\\boldsymbol{R}_{0}\\boldsymbol{D}_{\\boldsymbol{m}^{F}},\\\\\n",
    "\\boldsymbol{M}_{6}&:=\\left[\\begin{array}{ccc}\n",
    "\\boldsymbol{L}_{1} & \\boldsymbol{L}_{2} & \\boldsymbol{1}_{N_{tot}}\\end{array}\\right]^{\\top}\\boldsymbol{R}_{1},\n",
    "\\end{align*}\n",
    "\n",
    "which are represented by the matrices `R0` (of size $N_{tot}\\times F$, recall that $\\boldsymbol{1}_{N_{tot}^{f}}\\in\\mathbb{R}^{N_{tot}^{f}}$ is a $N_{tot}^{f}$-vector of $1$), `R1` (of size $N_{tot}\\times F$), and `M6` (of size $3\\times F$).\n",
    "\n",
    "### Lagrange multipliers\n",
    "\n",
    "Using the above definitions, we can compute the various quantities, which are:\n",
    "* the Pareto weights of the SWF: vector $\\boldsymbol{\\omega}$ and variable `ω`,\n",
    "* the marginal value of public fund $\\psi$: vector $\\boldsymbol{\\psi}$ and variable `ψ`,\n",
    "* the Lagrange multiplier for the Euler equation $(\\lambda_{h})_h$: vector $\\boldsymbol{\\lambda}$ and variable `λc`,\n",
    "* the one for the governmental budget constraint $\\mu$ (scalar `μ`),\n",
    "\n",
    "as well as the associated variables (quantities at the history level instead of the individual one): $\\bar{\\boldsymbol{\\omega}} :=\\boldsymbol{S} \\circ \\boldsymbol{\\omega}$ (variable `ωb`), $\\bar{\\boldsymbol{\\psi}} :=\\boldsymbol{S} \\circ \\boldsymbol{\\psi}$ (variable `ψb`), $\\bar{\\boldsymbol{\\lambda}} :=\\boldsymbol{S}\\circ\\boldsymbol{\\lambda}$ (variable `λcb`). Additionally, there is the average past Lagrange multiplier for the Euler equation: $(\\tilde\\lambda_{h})_h$: vector $\\tilde{\\boldsymbol{\\lambda}}$ and variable `λct`.\n",
    "\n",
    "We have the following relationships:\n",
    "\\begin{align*}\n",
    "    \\boldsymbol{S}\\circ\\tilde{\\boldsymbol{\\lambda}} &=\\boldsymbol{\\Pi}\\bar{\\boldsymbol{\\lambda}},\\\\\n",
    "    \\mu&=\\boldsymbol{L}_{0}^{\\top}\\bar{\\boldsymbol{\\omega}},\\\\\n",
    "    \\bar{\\boldsymbol{\\psi}}&=\\boldsymbol{M}_{4}\\bar{\\boldsymbol{\\omega}}\\\\\n",
    "    \\boldsymbol{\\bar{\\mathbf{\\lambda}}}&=\\boldsymbol{M}_{5}\\bar{\\boldsymbol{\\omega}}.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7b69eb0",
   "metadata": {},
   "source": [
    "## The weights determination  <a id=\"Ramsey-weights\"></a>[<font size=1>(back to Ramsey)</font>](#Ramsey)\n",
    "\n",
    "\n",
    "The Ramsey  FOCs deliver two additional constraints: \n",
    "\\begin{align*}\n",
    "    \\boldsymbol{L}_{1}\\boldsymbol{\\bar{\\omega}} &= 0, \\\\\n",
    "    \\boldsymbol{L}_{2}\\boldsymbol{\\bar{\\omega}} &= 0,\n",
    "\\end{align*}\n",
    "which are complemented by the normalization of social weights to 1: $\\boldsymbol{1}_{N_{tot}}\\bar{\\boldsymbol{\\omega}}=1$.\n",
    "\n",
    "These three equations yield: $\\boldsymbol{M}_{6}\\boldsymbol{\\omega}^{F}=\\left[\\begin{array}{ccc}\n",
    "0 & 0 & 1\\end{array}\\right]^{\\top}$. Since  $\\boldsymbol{R}_{1}$ is a $3\\times F$-matrix,   $\\boldsymbol{M}_{6}$ is a square matrix when $F=3$. The matrix  $\\boldsymbol{M}_{6}$  is in general invertible. We deduce:\n",
    "$$\\boldsymbol{\\omega}^{F}=\\boldsymbol{M}_{6}^{-1}\\left[\\begin{array}{ccc}0 & 0 & 1\\end{array}\\right]^{\\top}.$$\n",
    "\n",
    "Note that $\\boldsymbol{\\omega}^{F}$ is a vector of length $3$ representing the weights of each type. \n",
    "\n",
    "The vector of wieghts per history is $\\boldsymbol{\\omega}$, while $\\bar{\\boldsymbol{\\omega}}$ is the vector of weights per history multiplied by history sizes: $\\bar{\\boldsymbol{\\omega}}=(\\bar{\\omega}_h)_h=(S_h\\times \\omega_h)_h$. With our previous notation, $\\bar{\\boldsymbol{\\omega}}=\\boldsymbol{R}_{1}\\boldsymbol{\\omega}^{F}$ and:\n",
    "$$\\bar{\\boldsymbol{\\omega}}=\\boldsymbol{R}_{1}\\left[\\begin{array}{ccc}0 & 0 & 1\\end{array}\\right]^{\\top}.$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47097774",
   "metadata": {},
   "source": [
    "## The implementation  <a id=\"Ramsey-implementation\"></a>[<font size=1>(back to Ramsey)</font>](#Ramsey)\n",
    "\n",
    "The computation of the Ramsey allocation `ramsey_Solution::Ramsey` can be performed as follows:\n",
    "\n",
    "> `ramsey_Solution = Ramsey(truncatedModel::TruncatedModel,solution::AiyagariSolution,economy::Economy)`\n",
    "\n",
    "through a specific constructor of the struct `Ramsey`. \n",
    "\n",
    "This function takes as inputs: \n",
    "* economy parameters `economy::Economy`,\n",
    "* the solution of the Aiyagari model `solution::AiyagariSolution` corresponding to `economy`, \n",
    "* the truncated allocation `truncatedModel::TruncatedModel` of `solution`.\n",
    "\n",
    "The struct `ramsey_Solution` contains the following elements.\n",
    "* `weights::Weights` containing the SWF weights, and more precisely the vector of weights  $\\boldsymbol{\\omega}$ with length $|\\mathcal Y|$ (variable `ω`), and their transformation $\\overline{\\boldsymbol{\\omega}}$ (variable `ωb`). \n",
    "* `lagrangeMult::LagrangeMult` containing the steady-state values of the Lagrange multipliers of the Ramsey program used for  the Dynare simulation. \n",
    "* The truncated model `truncatedModel::TruncatedModel`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0cb7a4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "function get_weights(economy;N=2,R=2)\n",
    "    solution = steady(economy, print_step=20000, tolEGM=1e-12);\n",
    "    check_solution(solution,economy,noprint=true)\n",
    "    \n",
    "    refiNs = fill(N+R,maximum(economy.nys)) \n",
    "    truncatedModel = TruncatedModel(N,refiNs,solution,economy,method=\"closest\");\n",
    "    println(\"Number of constrained truncated histories: \",length(truncatedModel.truncatedAllocation.ind_cc_h))\n",
    "    check_truncation(truncatedModel,solution,economy,noprint=true)\n",
    "    return ramsey_GHH(truncatedModel,solution,economy)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "69214fdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "function Ramsey(truncatedModel::TruncatedModel,solution::AiyagariSolution,economy::Economy)#::Ramsey\n",
    "    @unpack N,Ntot,ind_h,ind_hs,truncatedAllocation,ξs = truncatedModel\n",
    "    @unpack S_h,Π_h,y0_h,a_beg_h,a_end_h,c_h,l_h,ly_τ_h,u_h,u′_h,u′′_h,v_h,v′_h,nb_cc_h,ind_cc_h = (\n",
    "                truncatedAllocation)\n",
    "    @unpack ξuE = ξs \n",
    "    @unpack β,α,δ,τ,u′,u′′,ys,ny,Sy,χ,φ,τ,nys= economy\n",
    "    @unpack w,R,K,L,G,Y,A = solution\n",
    "\n",
    "    T = typeof(w)\n",
    "    FL = (1-α)*K^α*L^-α\n",
    "\n",
    "    # be careful Π_h is foreward looking.\n",
    "    # S_h.*a_beg_h = Π_h'*(S_h.*a_end_h)\n",
    "    # S_h = Π_h'*(S_h)\n",
    "\n",
    "    ####\n",
    "    ## Preliminary definitions\n",
    "    ####\n",
    "\n",
    "    # Defining some scalars\n",
    "    l          = (χ*τ*w)^(1/(1.0/φ + 1.0 - τ))             # labor supply for a productivity of 1\n",
    "    τt         = (( 1.0/φ + 1.0)*τ)/( 1.0/φ +1.0 - τ )    # transformation of the progressivity \n",
    "    r          = R-1                                       # net interest rate\n",
    " \n",
    "    # Defining some specific matrices (see above for definitions)\n",
    "    Π_h_bar   = spdiagm(S_h)*Π_h*spdiagm(one(T)./S_h) #### verify ∀x :  Π_h_bar*(S_h.*x)=S_h.*(Π_h*x)  \n",
    "    Is        = sparse(I,Ntot,Ntot)\n",
    "    P         = Is - sparse(ind_cc_h,ind_cc_h,ones(T,nb_cc_h),Ntot,Ntot)   \n",
    "\n",
    "    M1        =  P* (Is - β*(1+r)*Π_h_bar)*spdiagm(ξuE.*u′′.(c_h,l_h))*(Is - (1+r)*Π_h') + (Is - P)\n",
    "    M2        = inv(Matrix(M1))*(P*(Is - β*(1+r)*Π_h_bar)*spdiagm(ξuE.*u′.(c_h,l_h)))\n",
    "    V1        = - M1 \\ (P*(Is - β*(1+r)*Π_h_bar)*S_h)\n",
    "\n",
    "    M3        = spdiagm(ξuE.*u′.(c_h,l_h))  - spdiagm(ξuE.*u′′.(c_h,l_h))* (Is - (1+r)*Π_h')*M2\n",
    "    V2        = -(spdiagm(ξuE.*u′′.(c_h,l_h))*(Is - (1+r)*Π_h')*V1 + S_h)\n",
    "\n",
    "    C1 = S_h'*( (l^(1.0/φ)/χ)*y0_h.^τt - FL*y0_h.^(1.0 +τt/(1.0/φ +1.0) )) - (1/(χ*τt))*(1.0/φ +1.0)*(l^(1.0/φ))*(y0_h.^τt)'*V2\n",
    "    L0 = ( (1/C1)*(1/(χ*τt))*(1.0/φ +1.0)*(l^(1.0/φ))*(y0_h.^τt)'*M3)'\n",
    "    \n",
    "    M4 = M3 + V2*L0'\n",
    "    M5 = M2 + V1*L0'\n",
    "\n",
    "    L1 = (a_beg_h'*M4 + (ξuE.*u′.(c_h,l_h))'*Π_h'*M5)' \n",
    "    L2 = (( (l^(1.0/φ+1.0)/(χ*τt))* ( (y0_h.^τt).*(log.(y0_h) - ones(length(c_h),1)/τt)) )'*M4 -\n",
    "            (l/(1.0/φ +1.0))*(S_h.*log.(y0_h))'*((l^(1.0/φ)/χ)*y0_h.^τt - FL*y0_h.^(1.0 +τt/(1.0/φ +1.0)) )*L0')'\n",
    "\n",
    "    Rs = Array{Matrix{T},1}(undef,length(ind_hs))\n",
    "    n_ind_hs = map(length,ind_hs)\n",
    "    for i in eachindex(n_ind_hs)\n",
    "        i0 = sum(n_ind_hs[1:i-1])\n",
    "        i1 = sum(n_ind_hs[1:i])\n",
    "        Rs[i] = [zeros(i0,1);S_h[i0+1:i1];zeros(Ntot-i1,1)]\n",
    "    end\n",
    "    R1  = hcat(Rs...)\n",
    "    M6 = vcat(L1'*R1,L2'*R1,sum(R1,dims=1))\n",
    "\n",
    "    # by bins\n",
    "    ω = M6 \\ [0;0;1]\n",
    "    ωb          = R1*ω\n",
    "    weights = Weights(ω,ωb,ωb,ωb,ωb,ωb,ωb,ωb) # the constructor allows for other wieghts that are not needed here\n",
    "\n",
    "    λc   = M5*ωb\n",
    "    λct  = Π_h'*λc \n",
    "    λcb  = M5*ωb\n",
    "    λl   = zeros(T,size(λc))   # not needed in this setup\n",
    "    λlt  = zeros(T,size(λc))   # not needed in this setup\n",
    "    λlb  = zeros(T,size(λc))   # not needed in this setup\n",
    "    μ    = L0'*ωb\n",
    "    ψ    = M4*ωb       # this is psi_hat\n",
    "    ψb   = M4*ωb       # this is psi_hat\n",
    "\n",
    "    lagrangeMult = LagrangeMult(λc,λct,λcb,λl,λlt,λlb,μ,ψ,ψb)\n",
    "\n",
    "return Ramsey(weights,lagrangeMult,truncatedModel,l)\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d861f0f5",
   "metadata": {},
   "source": [
    "### Checking the solution\n",
    "\n",
    "We check below different aspects of the Ramsey solution.\n",
    "* we check that $\\boldsymbol{S}$ is a eigenvector $\\boldsymbol{\\Pi}$ with eigenvalue 1;\n",
    "* we check that the definitions of $\\tilde{\\boldsymbol{a}}$ and $\\boldsymbol{a}$ are consistent with each other;\n",
    "* we check that the budget constraint holds at the history level;\n",
    "* we chack that the FOCs hold (wrt $a_h$, $l_h$, $R$, $w$, $\\tau$ in this order);\n",
    "* we check that $\\psi_h$ is null if $h$ is credit-constrained;\n",
    "* we check that the definition of $\\bar{\\boldsymbol{\\psi}}$ holds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8f9ed23",
   "metadata": {},
   "outputs": [],
   "source": [
    "function check_Ramsey(ramsey::Ramsey,solution::AiyagariSolution,economy::Economy;\n",
    "        noprint=false)::Nothing\n",
    "\n",
    "    @unpack weights,lagrangeMult,truncatedModel = ramsey\n",
    "    @unpack N,Ntot,ind_h,truncatedAllocation,ξs = truncatedModel\n",
    "    @unpack S_h,Π_h,y0_h,a_beg_h,a_end_h,c_h,l_h,ly_τ_h,u_h,u′_h,u′′_h,v_h,v′_h,nb_cc_h,ind_cc_h = (\n",
    "                truncatedAllocation)\n",
    "    @unpack ξuE = ξs \n",
    "    @unpack β,α,χ,δ,φ,τ,u′,u′′,ys,ny = economy\n",
    "    @unpack w,R,K,L,G,Y,A = solution\n",
    "    @unpack λcb,ψb,λlb,λc,λct,λl,λlt,ψ,μ = lagrangeMult\n",
    "    @unpack ωb_h,ω_h = weights\n",
    "    \n",
    "    # defining two scalars (labor supply l and transform of the progressivity)\n",
    "    l = (χ*τ*w)^(1/(1/φ+1-τ))\n",
    "    τt =  ((1+1/φ)*τ)/(1/φ+1-τ)\n",
    "\n",
    "    ψh = ψb\n",
    "    λ = λc    \n",
    "    FL = (1-α)*K^α*L^-α\n",
    "    Is        = sparse(I,Ntot,Ntot)\n",
    "    P         = Is - sparse(ind_cc_h,ind_cc_h,ones(typeof(β),nb_cc_h),Ntot,Ntot)   \n",
    "    Π_h_bar   = spdiagm(S_h)*Π_h*spdiagm(one(typeof(β))./S_h) #### verify ∀x :  Π_h_ψ*(S_h.*x)=S_h.(Π_h*x)  \n",
    "    log_ly_h  = log.(y0_h.*l_h)\n",
    "    \n",
    "    norm(x) = maximum(abs.(x))\n",
    "    diffs    = S_h - Π_h'*S_h\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in S_h. Diff (S_h - Π_h'*S_hb): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: S_h. Diff ||S_h - Π_h'*S_h||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    diffs    = S_h.*a_beg_h - Π_h'*(S_h.*a_end_h)\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in a_beg_h. Diff (a_beg_h): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: a_beg_h. Diff ||a_beg_h||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    diffs    = c_h + a_end_h - R*a_beg_h - w*ly_τ_h\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @show diffs\n",
    "        @warn(\"error in budget const. Diff (budget const): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: budget const. Diff ||budget const||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    diffs    = P*ψh - β*R*P*Π_h_bar*ψh\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in FOC (a_h). Diff (P*ψb - β*R*P*Π_h_bar*ψb): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: FOC (a_h). Diff ||P*ψb - β*R*P*Π_h_bar*ψb||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    #diffs    = a_beg_h'*ψh + (ξu1.*u′.(c_h,l_h))'*Π_h'*λ\n",
    "    diffs    = a_beg_h'*ψh + (ξuE.*u′.(c_h,l_h))'*Π_h'*λ\n",
    "    \n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in FOC (R). Diff (FOC): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: FOC (R). Diff ||FOC||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    diffs    = (l^(1/φ)*(1+1/φ)/(χ*τt))*(y0_h.^τt)'*ψh - μ*S_h'*(l^(1/φ)*y0_h.^τt/χ - FL*y0_h.^(1+τt*φ/(1+φ)))\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in FOC (w). Diff (FOC): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: FOC (w). Diff ||FOC||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    diffs    = l^(1/φ+1)/(χ*τt)*((log.(y0_h).- 1/τt).*y0_h.^τt)'*ψh - μ*l/(1/φ+1)*(\n",
    "            (S_h.*log.(y0_h))'*(l^(1/φ)*y0_h.^τt/χ - FL*y0_h.^(1+τt*φ/(1+φ))))\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in FOC (τ). Diff (FOC): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: FOC (τ). Diff ||FOC||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    diffs    = (Is-P)*λ\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in CC (λbc). Diff (CC): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: CC (λbc). Diff ||CC||∞: \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "\n",
    "    diffs = -ψh + ω_h.*ξuE.*u′.(c_h,l_h) - μ*S_h -(ξuE.*u′′.(c_h,l_h)).*((I - R*Π_h')*λ)\n",
    "    if !(norm(diffs) < 1e-10)\n",
    "        @warn(\"error in ψb. Diff (ψb - def_ψ): \", round.(diffs,digits=4))\n",
    "    else\n",
    "        (!noprint)&&println(\"Passed: ψb. Diff (ψb - def_ψ): \", round.(norm(diffs),digits=4))\n",
    "    end\n",
    "    \n",
    "    return nothing\n",
    "end;"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.5",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
